"""
sklearn对TF-IDF的优化：
TF-IDF = TF * IDF
TF: 直接使用词频统计 -> 用CountVectorizer 计算
IDF: 逆文档频率统计
    IDF = log(总文档数量 + 1  / 该词出现的文档数量 + 1) + 1
此外，还会进行L2归一化
"""
import numpy as np
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.preprocessing import normalize


def my_tfidf(data):
    # 1.计算词频 TF
    count_vec = CountVectorizer()
    tf = count_vec.fit_transform(data).toarray()    # 返回的是三元组，需要转换为ndarray
    print("词频矩阵：\n", tf)
    """
    把词频矩阵（稀疏矩阵，三元组）转换为数组（二维数组）
    1) 数组的每一行，就是一个文档 -> 数组元素数量 = 总文档数量
    2) 每一列的和 = 该词在所有文档中出现的次数，不能直接求和，因为有的文档中该词可能出现多次
        所以，要算 该词出现的文档数量，可以用布尔索引 （True=1，False=0）
    3） CountVectorizer 得到的，就是tf：该词在该文档出现的次数
    """
    corpus_count = len(tf)    # 总文档数量
    term_count = np.sum(tf!=0, axis=0)    # 包含该词的文档数量
    print("总文档数量：", corpus_count)
    print("包含该词的文档数量：", term_count)

    # 2. 计算逆文档频率 IDF
    idf = np.log((corpus_count + 1) / (term_count + 1)) + 1
    print("逆文档频率：", idf)
    # 3. 计算TF-IDF
    tf_idf = tf * idf
    # 4. 归一化
    tf_idf_norm = normalize(tf_idf, norm="l2")
    print("TF-IDF矩阵：\n", tf_idf_norm)

    return None


if __name__ == '__main__':
    corpus = ['I love programming',
              'Programming is fun']
    my_tfidf(corpus)